from exper.Experiment import Experiment
import torch
import os
import math
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import scipy
import matplotlib

class ExpUntarget(Experiment):

    def __init__(self, config, dataer, neter):

        super(ExpUntarget, self).__init__('Untarget', config)
        self.dataer = dataer
        self.neter = neter
    
    def latex_output(self, case='Case1'):
        
        ## 转置成DFdata吧，然后存起来作为实验放图，然后再转它 

        dicter = self.load_file(case=case)
        if case == 'Case1':
            list = ['OnePixel', 'SimBA', 'FNS', 'MPack', 'GA', 'SMPack', 'FGSM']
            str = ''
            for method in list:
                str += '&${:.1f}\\pm{:.1f}$ '.format(dicter[method][0] * 100, math.sqrt(dicter[method][1] * 100))
            str += '\\\\'
        elif case == 'Case2':
            list = ['OnePixel', 'SimBA', 'FNS', 'MPack', 'GA', 'SMPack']
            str = ''
            for method in list:
                str += '&${:.1f}\\pm{:.1f}$ \n'.format(dicter[method]['Average'][0], math.sqrt(dicter[method]['Average'][1]))
            str += '\n\n\n'
            for method in list:
                str += '&${:.1f}\\pm{:.1f}$ \n'.format(dicter[method]['Median'][0], math.sqrt(dicter[method]['Median'][1]))
            str += '\\\\'
        else:
            raise ValueError
        
        return str

    def Case_1_2(self, seed, repeat_times=3, default_num=500, way='SMPack', atk=None):

        ASR_list = []
        avg_Query_list = []
        median_Query_list = []

        for step in range(repeat_times):

            loader = self.dataer.get_shuffle_loader(seed=seed+step, isTrain=False, batch_size=1)
            succ_number = 0
            query_buffer = 0
            single_sample_query_list = []

            count = 0
            total = 0

            for (image, label) in loader:
                
                if count == default_num:
                    break
                total += 1
                test_image = image.to(self.neter.device)
                test_label = label.to(self.neter.device)
                output = self.neter.net(test_image)
                _, pre = torch.max(output.data, 1)
                if pre != test_label:
                    print(pre)
                    print(test_label)
                    print('Pass the error sample')
                    continue

                if way == 'SMPack':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.untarget_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                        batch=self.config['base_batch'],
                                        size=self.config['base_size'],
                                        mode=self.config['base_mode'],
                                    )
                    if flag:
                        succ_number += 1         
                
                elif way == 'GA':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.untarget_GA_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                    )
                    if flag:
                        succ_number += 1     
                   
                elif way == 'MPack':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.untarget_MPack_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                        batch=self.config['base_batch'],
                                        size=self.config['base_size'],
                                        mode=self.config['base_mode'],
                                    )
                    if flag:
                        succ_number += 1   

                elif way == 'FNS':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.FNS_attack(
                                        images=image,
                                        labels=label,
                                    )
                    if flag:
                        succ_number += 1 

                elif way == 'OnePixel':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1   

                elif way == 'SimBA':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.SimBA_attack(image, label[0], max_queries=10000, epsilon=self.config['epsilon'], targeted=False)

                    if flag:
                        succ_number += 1
                        
                elif way == 'FGSM':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1  
                
                elif way == 'PGD':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1  

                elif way == 'ZOO':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image, _ = atk.attack(image, label)
                    adv_image = torch.from_numpy(adv_image).unsqueeze(0)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1  

                elif way == 'Test':
                    pass
                else:
                    raise Exception('No such compared method.')
                print('Total : {}'.format(total))
                ## query counts
                print('='*50)
                print('Qurey for {}-th samples : {}'.format(count, self.neter.query - query_buffer))
                count += 1
                print('Total Query: {}'.format(self.neter.query))
                print('Last Total Query: {}'.format(query_buffer))
                print('='*50)
                single_sample_query_list.append(self.neter.query - query_buffer)
                query_buffer = self.neter.query

            print('Way {} : The Epoch {}-th Untarget Attack ASR (Randomly sample {} samples) : {:.2f}%'.format(way, step, default_num, succ_number / default_num * 100))
            
            ASR_list.append(succ_number / default_num) 
            avg_Query_list.append(self.neter.query / default_num)
            
            # median
            single_sample_query_list.sort()
            list_len = len(single_sample_query_list)
            if list_len % 2 == 0:
                median_Query_list.append((single_sample_query_list[int(list_len / 2) - 1] + single_sample_query_list[int(list_len / 2)]) / 2)
            else:
                median_Query_list.append(single_sample_query_list[int(list_len / 2)])

            self.neter.clear_query()

        case_1_path = self.generate_name(case='Case1')
        case_2_path = self.generate_name(case='Case2')

        # for case 1:
        if os.path.exists(case_1_path) == False:
            os.makedirs(case_1_path)
        if os.path.exists(os.path.join(case_1_path, 'result.pt')) == False:
            case1_dicter = {}
            case1_dicter[way] = [np.mean(ASR_list), np.var(ASR_list)]
        else:
            case1_dicter = torch.load(os.path.join(case_1_path, 'result.pt'))
            case1_dicter[way] = [np.mean(ASR_list), np.var(ASR_list)]
        torch.save(case1_dicter, os.path.join(case_1_path, 'result.pt'))

        # for case 2:
        if os.path.exists(case_2_path) == False:
            os.makedirs(case_2_path)
        if os.path.exists(os.path.join(case_2_path, 'result.pt')) == False:
            case2_dicter = {}
            case2_dicter[way] = {
                'Average': [np.mean(avg_Query_list), np.var(avg_Query_list)],
                'Median': [np.mean(median_Query_list), np.var(median_Query_list)],
                }
        else:
            case2_dicter = torch.load(os.path.join(case_2_path, 'result.pt'))
            case2_dicter[way] = {
                'Average': [np.mean(avg_Query_list), np.var(avg_Query_list)],
                'Median': [np.mean(median_Query_list), np.var(median_Query_list)],
            }
        torch.save(case2_dicter, os.path.join(case_2_path, 'result.pt'))

    def Case_3(self, atk, rows=3):

        """
        绘制 原始图片 MPack SMPack 三列图片
        选取三个好看的case作案例 
        """
        classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
        
        sns.set_style('white')
        fig = plt.figure(figsize=(9, 10))

        loader = self.dataer.get_shuffle_loader(seed=0, isTrain=False, batch_size=1)

        index = 1
        count = 0

        succ_class_list = []

        for row in range(rows):
            
            while True:

                for dex, (image, label) in enumerate(loader):
                    
                    if dex == count:
                        count += 1
                        break

                test_image = image.to(self.neter.device)
                test_label = label.to(self.neter.device)

                output = self.neter.net(test_image)
                _, pre = torch.max(output.data, 1)
                if pre != test_label or pre[0] in succ_class_list:
                    print(pre)
                    print(test_label)
                    print('Pass the error or class same sample')
                    continue

                ## GA
                flag2, adv_image2 = atk.untarget_GA_attack(
                                origin_sample=image,
                                source_class=label[0],
                                epsilon=self.config['epsilon'],
                                norm_type=self.config['attck_norm'],
                                method=self.config['attck_method'],
                            )
                cuda_adv_image2 = adv_image2.to(self.neter.device)
                output = self.neter.net(cuda_adv_image2)
                _, label2 = torch.max(output.data, 1)


                ## SMPack 
                flag3, adv_image3 = atk.untarget_attack(
                                    origin_sample=image,
                                    source_class=label[0],
                                    epsilon=self.config['epsilon'],
                                    norm_type=self.config['attck_norm'],
                                    method=self.config['attck_method'],
                                    batch=self.config['base_batch'],
                                    size=self.config['base_size'],
                                    mode=self.config['base_mode'],
                                )
                cuda_adv_image3 = adv_image3.to(self.neter.device)
                output = self.neter.net(cuda_adv_image3)
                _, label3 = torch.max(output.data, 1)

                if flag2 == True and flag3 == True:

                    succ_class_list.append(label[0])

                    break
            
            numpy_img = image.squeeze().numpy()
            numpy_adv_image2 = adv_image2.squeeze().numpy()
            numpy_adv_image3 = adv_image3.squeeze().numpy()
            
            numpy_img =  np.transpose(numpy_img, (2, 1, 0))
            numpy_adv_image2 =  np.transpose(numpy_adv_image2, (2, 1, 0))
            numpy_adv_image3 =  np.transpose(numpy_adv_image3, (2, 1, 0))
            
            
            ## origin sample
            ax1 = plt.subplot(rows, 3, index)
            index += 1
            tr_numpy_img = scipy.ndimage.rotate(numpy_img, -90)
            ax1.imshow(tr_numpy_img)
            ax1.set_xticks([])
            ax1.set_yticks([])
            ax1.set_xlabel('{}'.format(classes[label[0]]), fontsize=16)
            if index == 2:
                ax1.set_title('Origin', fontsize=14)

            ## GA
            ax2 = plt.subplot(rows, 3, index)
            index += 1
            tr_numpy_adv_image2 = scipy.ndimage.rotate(numpy_adv_image2, -90)
            ax2.imshow(tr_numpy_adv_image2)
            ax2.set_xticks([])
            ax2.set_yticks([])
            ax2.set_xlabel('{}'.format(classes[label2[0]]), fontsize=16)
            if index == 3:
                ax2.set_title('GA', fontsize=14)

            ## SMPack
            ax3 = plt.subplot(rows, 3, index)
            index += 1
            tr_numpy_adv_image3 = scipy.ndimage.rotate(numpy_adv_image3, -90)
            ax3.imshow(tr_numpy_adv_image3)
            ax3.set_xticks([])
            ax3.set_yticks([])
            ax3.set_xlabel('{}'.format(classes[label3[0]]), fontsize=16)
            if index == 4:
                ax3.set_title('$\it{SMPack}$', fontsize=14)

        plt.subplots_adjust(wspace=0.2, hspace=0.3)
        plt.savefig('./img/untarget_show.pdf', dpi=600, bbox_inches='tight', pad_inches=0.2)
        print('Save done')
        plt.close()
            
    def norm_calculate(self, seed=0, default_num=10, atk=None):

        loader = self.dataer.get_shuffle_loader(seed=seed, isTrain=False, batch_size=1)
        
        count = 0
        total = 0

        norm1 = []
        norm2 = []
        norm3 = []

        for (image, label) in loader:

            print(count)
            
            if count == default_num:
                break

            total += 1

            test_image = image.to(self.neter.device)
            test_label = label.to(self.neter.device)
            output = self.neter.net(test_image)
            _, pre = torch.max(output.data, 1)
            if pre != test_label:
                print(pre)
                print(test_label)
                print('Pass the error sample')
                continue

            flag1, adv_image1 = atk.untarget_attack(
                                origin_sample=image,
                                source_class=label[0],
                                epsilon=self.config['epsilon'],
                                norm_type=self.config['attck_norm'],
                                method=self.config['attck_method'],
                                batch=self.config['base_batch'],
                                size=self.config['base_size'],
                                mode=self.config['base_mode'],
                            )
   
            flag2, adv_image2 = atk.untarget_GA_attack(
                                origin_sample=image,
                                source_class=label[0],
                                epsilon=self.config['epsilon'],
                                norm_type=self.config['attck_norm'],
                                method=self.config['attck_method'],
                            )
  
            flag3, adv_image3 = atk.untarget_MPack_attack(
                                origin_sample=image,
                                source_class=label[0],
                                epsilon=self.config['epsilon'],
                                norm_type=self.config['attck_norm'],
                                method=self.config['attck_method'],
                                batch=self.config['base_batch'],
                                size=self.config['base_size'],
                                mode=self.config['base_mode'],
                            )
            
            if flag1 == True and flag2 == True and flag3 == True:
                
                norm1.append(torch.norm((adv_image1 - image).reshape(-1)).detach().cpu().numpy())
                norm2.append(torch.norm((adv_image2 - image).reshape(-1)).detach().cpu().numpy())
                norm3.append(torch.norm((adv_image3 - image).reshape(-1)).detach().cpu().numpy())
                count += 1
            
            else:
                
                continue
        
        if len(norm1) != default_num:
            raise Exception('No eq default number')
        

        print('{} {}'.format(self.config['model_arch'], self.config['attck_method']))
        print('&${}\pm{}$ &${}\pm{}$ &${}\pm{}$ \\\\'.format(np.mean(norm3), np.var(norm3), np.mean(norm2), np.var(norm2), np.mean(norm1), np.var(norm1)))

    def build_matrix(self, default_num=100, atk=None, way='MPack'):


            matrix = np.zeros((self.dataer.class_num, self.dataer.class_num))

            source_classes = [class_index for class_index in range(self.dataer.class_num)]

            for source_class in source_classes:
                    
                count = 0
                total = 0                

                loader = self.dataer.get_class_loader(class_index=source_class, isTrain=False, batch_size=1)

                for (image, label) in loader:
                    
                    if count == default_num:
                        break
                    total += 1
                    test_image = image.to(self.neter.device)
                    test_label = label.to(self.neter.device)
                    output = self.neter.net(test_image)
                    _, pre = torch.max(output.data, 1)
                    if pre != test_label:
                        print(pre)
                        print(test_label)
                        print('Pass the error sample')
                        continue

                    print('Attacking {}-th sample'.format(count))

                    if way == 'SMPack':

                        flag, adv_image = atk.untarget_attack(
                                            origin_sample=image,
                                            source_class=label[0],
                                            epsilon=self.config['epsilon'],
                                            norm_type=self.config['attck_norm'],
                                            method=self.config['attck_method'],
                                            batch=self.config['base_batch'],
                                            size=self.config['base_size'],
                                            mode=self.config['base_mode'],
                                        )
                    elif way == 'MPack':
                    
                        flag, adv_image = atk.untarget_MPack_attack(
                                            origin_sample=image,
                                            source_class=label[0],
                                            epsilon=self.config['epsilon'],
                                            norm_type=self.config['attck_norm'],
                                            method=self.config['attck_method'],
                                            batch=self.config['base_batch'],
                                            size=self.config['base_size'],
                                            mode=self.config['base_mode'],
                                        )
                        
                    if adv_image == None:
                        adv_image = image

                    adv_image = adv_image.to(self.neter.device)
                    adv_output = self.neter.net(adv_image)
                    _, pre = torch.max(adv_output.data, 1)

                    matrix[source_class][pre[0]] += 1
                    count += 1
                
            np.save(os.path.join(self.basic_path, '{}-matrix.npy'.format(way)), matrix)
            print('Save done')

    def plot_matrix(self, default_num=100, atk=None, way='MPack'):

        if os.path.exists(os.path.join(self.basic_path, '{}-matrix.npy'.format(way))) == False:

            print('Do not find matrix data')
            self.build_matrix(default_num=default_num, atk=atk, way=way)
        
        else:

            print('Find matrix data')

        matrix = np.load(os.path.join(self.basic_path, '{}-matrix.npy'.format(way)))
        print('Load done')
        print(matrix)

        ## tick_labels = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
        
        ax = sns.heatmap(
            matrix, 
            annot=True, 
            cmap='Blues', 
            cbar=False, 
            linewidths=.5, 
            square=True, 
            )
        # ax.xaxis.tick_top()
        ## ax.set(xlabel='Classified Label', ylabel='Source Label')
        ax.set_xlabel(xlabel='Classified Label', fontsize=14)
        ax.set_ylabel(ylabel='Source Label', fontsize=14)
        # ax.xaxis.set_label_position('top') 
        if way == 'MPack':
            plt.title('Untarget Attack   VGG-16   MPack', fontsize=16)
        elif way == 'SMPack':
            plt.title('Untarget Attack   VGG-16   SMPack', fontsize=16)
        ax.axhline(y=0, color='k', linewidth=.5)
        ax.axhline(y=10, color='k', linewidth=.5)
        ax.axvline(x=0, color='k', linewidth=.5)
        ax.axvline(x=10, color='k', linewidth=.5)
        
        plt.savefig('./img/Untarget_{}_matrix.pdf'.format(way), dpi=1200, bbox_inches='tight', pad_inches=0.2)
        
        print('Save done')